
#include <stack>
#include <boost/algorithm/string.hpp>
#include "stl-helper.h"
#include "vom.h"


VOM::VOM()
{
  root_ = new VOMNode;
}

void VOM::build_big_tree(VOMNode* node, const SequencesSA& seqsa)
{
  if (node->count_ < 2) return;
  if (node->parent_ == NULL) tmp_s = "";
  else tmp_s = node->c_ + tmp_s;

  std::string acgt = "ACGT";
  //MSG("tmp_s = " << tmp_s);
  BOOST_FOREACH(char c, acgt) {
    int count = seqsa.count(c+tmp_s);

    //MSG("tmp_s+c = " << tmp_s+c << " count=" << count);
    
    if (count >= 1) {
      VOMNode* new_node = new VOMNode(c, node);
      new_node->count_ = count;
      node->add_child(new_node);
    }
  }

  BOOST_FOREACH(VOMNode* child, node->child_node_) {
    build_big_tree(child, seqsa);
  }

  if (node->parent_ == NULL) tmp_s = "";
  else tmp_s.erase(0, 1);
}

void VOM::build_fix_tree(VOMNode* node, const SequencesSA& seqsa, int remain_depth)
{
  if (remain_depth == 0) return;

  if (node->parent_ == NULL) tmp_s = "";
  else tmp_s = node->c_ + tmp_s;

  std::string acgt = "ACGT";
  //MSG("tmp_s = " << tmp_s);
  BOOST_FOREACH(char c, acgt) {
    int count = seqsa.count(c+tmp_s);

    //MSG("tmp_s+c = " << tmp_s+c << " count=" << count);
    
    VOMNode* new_node = new VOMNode(c, node);
    new_node->count_ = count;
    node->add_child(new_node);

  }

  BOOST_FOREACH(VOMNode* child, node->child_node_) {
    build_fix_tree(child, seqsa, remain_depth-1);
  }

  if (node->parent_ == NULL) tmp_s = "";
  else tmp_s.erase(0, 1);
}

std::ostream& VOM::print(std::ostream& o, bool isFile) const
{
  print(o, root_, 0);
  return o;
}

std::ostream& VOM::print(std::ostream& o, VOMNode* node, int depth) const
{
  for (int i = 0; i < depth; i++) o << "  ";

  o << node->c_ << " (" << node->count_ << ")";
  for (uint i = 0; i < node->prob_.size(); i++)
    o << " " << node->prob_[i];
  o  << ENDL;
  BOOST_FOREACH(VOMNode* child, node->child_node_)
    print(o, child, depth+1);
  o  << "UP" << ENDL;
  return o;
}

void VOM::mle_estimate(VOMNode* node, VOMNode* na, VOMNode* nc, VOMNode* ng, VOMNode* nt)
{
  node->prob_ = arma::zeros<vd_t>(4);
  if (node->count_ > 0) {
    if (na != NULL) node->prob_[0] = (double) na->count_ / node->count_;
    if (nc != NULL) node->prob_[1] = (double) nc->count_ / node->count_;
    if (ng != NULL) node->prob_[2] = (double) ng->count_ / node->count_;
    if (nt != NULL) node->prob_[3] = (double) nt->count_ / node->count_;
  }

  double pseudo_count = 0.1;
  for (uint i = 0; i < node->prob_.size(); i++)
    if (node->prob_[i] >0 && node->prob_[i] < pseudo_count) pseudo_count = node->prob_[i];
  node->prob_ += pseudo_count / 10;
  node->prob_ /= arma::sum(node->prob_);

  BOOST_FOREACH(VOMNode* child, node->child_node_) {
    char c = child->c_;
    mle_estimate(child, na ? na->child(c) : NULL, nc ? nc->child(c) : NULL, 
	  ng ? ng->child(c) : NULL, nt ? nt->child(c) : NULL);
  }
}

double VOM::log_likelihood(const std::string& seq, int max_order) const
{
  double ll = 0;
  for (uint i = 0; i < seq.size(); i++) {
    const vd_t& p = prob_vec(seq, i-1, max_order < 0 ? i : max_order);
    double prob = p[seq[i] == 'A' ? 0 : (seq[i] == 'C' ? 1 : (seq[i] == 'G' ? 2 : 3))];
    ll += log(prob);
  }
  return ll;
}

const vd_t& VOM::prob_vec(const std::string& seq, int i, int max_order) const
{
  VOMNode* node = root_;
  int start = i;
  while (i >= 0 && i > start - max_order) {
    //MSG(start << " " << i);
    VOMNode* child = node->child(seq[i--]);
    if (!child) break;
    node = child;
  }
  return node->prob_;
}

void VOM::load(const std::string& modelFile)
{
  delete root_;
  root_ = new VOMNode;

  FileAccessRead(file, modelFile);
  std::stack<VOMNode*> s;
  s.push(root_);
  while (!file.eof() && !s.empty()) {
    std::string line;
    std::getline(file, line);

    std::vector<std::string> tokens;
    boost::split(tokens, line, boost::is_any_of("()\t "), boost::token_compress_on);
    
    if (sz(tokens) == 0) continue;
    int start = 0;
    if (tokens[0] == "") start = 1;
    
    // BOOST_FOREACH(std::string tok, tokens)
    //   std::cerr << tok << "-";
    // std::cerr << " size = " << tokens.size() 
    // 	      << " start = " << tokens[start] << ENDL;

    if (tokens[start] == "UP") {
      s.pop();
    }
    else if (tokens[start] == "X") {
      root_->parent_ = NULL;
      root_->c_ = 'X';
      root_->count_ = RAND::s2type<int>(tokens[start+1]);
      root_->prob_ = arma::zeros<vd_t>(4);
      for (int i = 0; i < 4; i++)
	root_->prob_[i] = RAND::s2type<double>(tokens[start+2+i]);
    }
    else {
      VOMNode* node = s.top();
      VOMNode* child = new VOMNode(tokens[start][0], node);
      
      child->count_ = RAND::s2type<int>(tokens[start+1]);
      child->prob_ = arma::zeros<vd_t>(4);
      for (int i = 0; i < 4; i++)
	child->prob_[i] = RAND::s2type<double>(tokens[start+2+i]);
      
      node->add_child(child);
      s.push(child);
    }     
    //std::cout << line;
  }
}
